from models import F3Net
import torch
import torch.nn as nn
from torch.nn import parameter
import torch.nn.functional as F
import numpy as np
import os
import cv2
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
osenvs = len(os.environ['CUDA_VISIBLE_DEVICES'].split(','))

class detect_Net():
    def __init__(self, gpu_ids, pretrained_path, mode ='Both'):
        self.device = torch.device('cuda:{}'.format(gpu_ids[0])) if gpu_ids else torch.device('cpu')
        self.net = F3Net(mode=mode, device=self.device)
        self.model = self.initModel(self.model, gpu_ids)
        self.load(pretrained_path)
        self.gpu_ids = gpu_ids

    def forward(self, x):
        fea, out = self.model(x)
        del fea
        return out

    def detect_image(self, image):

        #图片的预处理



        with torch.no_grad():
            image = torch.from_numpy(image)
            image = image.to(f'cuda:{self.gpu_ids[0]}')
            result = self.forward(image)
            #后处理

            print(result)


    def load(self, path):
        state_dict = torch.load(path)
        self.model.load_state_dict(state_dict)

    def initModel(self, mod, gpu_ids):
        mod = mod.to(f'cuda:{gpu_ids[0]}')
        mod = nn.DataParallel(mod, gpu_ids)
        return mod
    @staticmethod
    def set_input(self, input, label):
        self.input = input.to(self.device)
        self.label = label.to(self.device)



if __name__ == "__main__":
    image_path = 'dataset/'
    weight_path = 'pre_train_weight/xception-b5690688.pth'
    gpu_ids = [*range(osenvs)]
    detect = detect_Net(gpu_ids=gpu_ids, pretrained_path = weight_path, mode= 'Both')
    image = cv2.imread(image_path)

